import traceback
from datetime import datetime, timedelta

import pandas as pd
import pyqtgraph as pg
from PyQt5 import QtGui, QtCore
from PyQt5.QtGui import QIcon
from PyQt5.QtWidgets import QWidget
import akshare as ak
from pyfinance import TSeries

from forms.widgetGraph import Ui_FormGraph
from utils.BaseSignal import BaseSignal


class WidgetGraph(QWidget, Ui_FormGraph, BaseSignal):
    def __init__(self, parent=None):
        super(WidgetGraph, self).__init__(parent=parent)
        self.data_list = []
        self.bar_view = pg.ViewBox()
        self.hLine = pg.InfiniteLine(angle=0, movable=False, )  # 创建一个水平线条
        self.vLine = pg.InfiniteLine(angle=90, movable=False, )  # 创建一个垂直线条
        self.labelBlockText = pg.TextItem()  # 创建一个文本项
        self.k_plt = pg.PlotWidget()

        self.setupUi(self)

        # 当前股票代码
        self.code = '000000'
        # 当前股票名称
        self.name = ''

        self.data = pd.DataFrame()
        self.df_5d = pd.DataFrame()
        self.df_10d = pd.DataFrame()
        self.chg = TSeries()

        self.axis_dict = dict()

        self.move_slot = pg.SignalProxy(self.k_plt.scene().sigMouseMoved, rateLimit=60, slot=self.print_slot)

    def setupUi(self, FormGraph):
        super(WidgetGraph, self).setupUi(self)
        self.gridLayout.addWidget(self.k_plt)
        self.resize(1400, 700)
        self.k_plt.setLabel(axis='left', text='指数')  # 设置Y轴标签
        self.k_plt.setLabel(axis='bottom', text='日期')  # 设置X轴标签
        self.k_plt.showGrid(x=True, y=True)  # 设置绘图部件显示网格线
        self.setWindowIcon(QIcon("../icons/graph.png"))

    def show(self) -> None:
        super(WidgetGraph, self).show()

    @QtCore.pyqtSlot()
    def on_btnShowBarView_clicked(self):
        if self.btnShowBarView.text() == "显示涨跌柱状图":
            # 切换按钮显示文字
            self.btnShowBarView.setText("显示k线图")
            # 计算涨跌值数组
            updown = self.data['收盘'] - self.data['开盘']
            # 画柱状图，涨跌值小于零则画成绿色，否则是红色
            bar_chart = pg.BarGraphItem(x=range(0, len(self.data['日期'])),
                                        height=abs(updown),
                                        width=0.5, brushes=[QtGui.QColor('green') if updown[i] < 0
                                                            else QtGui.QColor('red') for i in range(len(updown))])
            # 清空绘图项目，重新画，注意这里不会清除坐标轴标签
            self.k_plt.plotItem.clear()
            self.k_plt.addItem(bar_chart)

            # 这里绘制方框文本信息
            self.k_plt.addItem(self.labelBlockText)  # 在图形部件中添加文本项

            # 这里绘制十字准线
            self.k_plt.addItem(self.vLine)  # 在图形部件中添加垂直线条
            self.k_plt.addItem(self.hLine)  # 在图形部件中添加水平线条

            self.k_plt.setYRange(0, max(abs(updown)))
        elif self.btnShowBarView.text() == "显示k线图":
            self.btnShowBarView.setText("显示涨跌柱状图")
            self.k_plt.plotItem.clear()

            y_min = self.data['最低'].min()
            y_max = self.data['最高'].max()
            self.k_plt.setYRange(y_min, y_max)
            item = CandlestickItem(self.data_list)  # 生成蜡烛图数据
            self.k_plt.addItem(item, )  # 在绘图部件中添加蜡烛图项目
            self.k_plt.addItem(self.labelBlockText)  # 在图形部件中添加文本项
            self.k_plt.addItem(self.vLine)  # 在图形部件中添加垂直线条
            self.k_plt.addItem(self.hLine)  # 在图形部件中添加水平线条
            # self.k_plt.addItem(pg.TextItem())
            self.k_plt.addLegend(size=(150, 80))
            self.k_plt.plot(y=self.df_5d.values, pen='#33a3dc', name='5日均线')
            self.k_plt.plot(y=self.df_10d.values, pen='#f173ac', name='10日均线')

    @QtCore.pyqtSlot()
    def on_btnQuery_clicked(self):
        days = 0
        if self.radioBtn1m.isChecked():
            days = 30
        elif self.radioBtn3m.isChecked():
            days = 90
        elif self.radioBtn6m.isChecked():
            days = 180
        elif self.radioBtn1y.isChecked():
            days = 365
        elif self.radioBtn3y.isChecked():
            days = 365 * 3
        self.query_slot(self.code, days)

    # 查询股票历史行情
    def query_slot(self, sym: str, days: int = 365):
        try:
            self.signalShowProgressBar.emit()
            self.signalUpdateProgressBar.emit(0)
            # 获取开始时间，即当前时间减去一年的天数
            start_date = datetime.today() - timedelta(days=days)
            start_date_str = datetime.strftime(start_date, "%Y%m%d")
            # 结束时间：今天
            end_date = datetime.today()
            end_date_str = datetime.strftime(end_date, "%Y%m%d")
            self.signalUpdateProgressBar.emit(30)
            self.labelStockName.setText(self.name)
            self.labelStockCode.setText(self.code)
            return self.plot_k_line(symbol=sym, start_date=start_date_str, end_date=end_date_str)
        except Exception:
            print(traceback.print_exc())

    def plot_k_line(self, symbol=None, start_date=None, end_date=None):
        self.data = ak.stock_zh_a_hist(symbol=symbol, period="daily",
                                       start_date=start_date, end_date=end_date,
                                       adjust="hfq").dropna()
        self.signalUpdateProgressBar.emit(40)
        if self.data.empty:
            return False
        self.data.index = pd.to_datetime(self.data['日期'])
        self.data.sort_index()

        self.btnShowBarView.setText("显示涨跌柱状图")
        # 涨跌幅
        self.chg = TSeries(self.data['收盘'] / self.data['收盘'].shift(1) - 1)
        # 年化收益率
        anl_ret = self.chg.anlzd_ret()
        # 累计收益率
        cum_ret = self.chg.cuml_ret()
        self.labelAulRate.setText(f'{anl_ret * 100:.2f}%')
        self.labelCumRate.setText(f'{cum_ret * 100:.2f}%')

        # benchmark = self.chg.loc[self.chg.index]
        ir = self.chg.sharpe_ratio()
        sor = self.chg.sortino_ratio(freq=250)
        cr = self.chg.calmar_ratio()
        self.label_sr.setText(f'{ir:.2f}')
        self.label_sor.setText(f'{sor:.2f}')
        self.label_cr.setText(f'{cr:.2f}')
        self.data_list.clear()
        y_min = self.data['最低'].min()
        y_max = self.data['最高'].max()
        d = 0
        for dates, row in self.data.iterrows():
            _open, close, high, low = row[1:5]
            datas = (d, _open, close, low, high)
            self.data_list.append(datas)
            d += 1
        self.axis_dict = dict(enumerate(self.data['日期']))
        date_list = list(self.data['日期'])
        axis_1 = [(i, date_list[i]) for i in range(0, len(date_list), 3)]  # 获取日期值
        axis_2 = [(i, date_list[i]) for i in range(0, len(date_list), 5)]
        axis_3 = [(i, date_list[i]) for i in range(0, len(date_list), 8)]
        axis_4 = [(i, date_list[i]) for i in range(0, len(date_list), 30)]
        axis_5 = [(i, date_list[i]) for i in range(0, len(date_list), 180)]
        self.signalUpdateProgressBar.emit(60)
        # stringaxis = pg.AxisItem(orientation='bottom')  # 创建一个刻度项
        # stringaxis.setTicks([axis_5, axis_4, axis_3, axis_2, axis_1, self.axis_dict.items()])  # 设置X轴刻度值
        self.k_plt.getAxis("bottom").setTicks([axis_5, axis_4, axis_3, axis_2, axis_1, self.axis_dict.items()])

        self.k_plt.setXRange(0, len(self.data['日期']))
        self.k_plt.setYRange(y_min, y_max)

        self.k_plt.plotItem.clear()  # 清空绘图部件中的项

        # 五日均线和十日均线数据
        self.df_5d = self.data['收盘'].rolling(5).mean()
        self.df_10d = self.data['收盘'].rolling(10).mean()

        item = CandlestickItem(self.data_list)  # 生成蜡烛图数据
        self.k_plt.addItem(item, )  # 在绘图部件中添加蜡烛图项目
        self.k_plt.addItem(self.labelBlockText)  # 在图形部件中添加文本项
        self.k_plt.addItem(self.vLine)  # 在图形部件中添加垂直线条
        self.k_plt.addItem(self.hLine)  # 在图形部件中添加水平线条

        self.k_plt.addLegend(size=(150, 80))
        self.k_plt.plot(y=self.df_5d.values, pen='#33a3dc', name='5日均线')
        self.k_plt.plot(y=self.df_10d.values, pen='#f173ac', name='10日均线')

        self.signalUpdateProgressBar.emit(100)
        self.signalFinished.emit()
        # 响应鼠标移动绘制十字光标
        return True

    def print_slot(self, event=None):
        if event is None:
            print("事件为空")
        else:
            pos = event[0]  # 获取事件的鼠标位置
            try:
                # 如果鼠标位置在绘图部件中
                if self.k_plt.sceneBoundingRect().contains(pos):
                    mouse_point = self.k_plt.plotItem.vb.mapSceneToView(pos)  # 转换鼠标坐标
                    index = int(mouse_point.x())  # 鼠标所处的X轴坐标
                    # pos_y = int(mouse_point.y())  # 鼠标所处的Y轴坐标
                    chg = self.chg.round(2).values
                    if 0 <= index < len(self.data.index):
                        # 在label中写入HTML
                        self.labelBlockText.setHtml(
                            "<p style='color:white'><strong>日期：{0}</strong></p>"
                            "<p style='color:white'>开盘：{1}</p>"
                            "<p style='color:white'>收盘：{2}</p>"
                            "<p style='color:white'>最高价：<span style='color:red;'>{3}</span></p>"
                            "<p style='color:white'>最低价：<span style='color:green;'>{4}</span></p>"
                            "<p style='color:white'>涨跌幅：<span style='color:{5};'>{6:.2f}%</span></p>".format(
                                self.axis_dict[index], self.data['开盘'][index], self.data['收盘'][index],
                                self.data['最高'][index], self.data['最低'][index], 'red' if chg[index] >= 0 else 'green',
                                chg[index] * 100
                            ))
                        self.labelBlockText.setPos(mouse_point.x(), mouse_point.y())  # 设置label的位置
                    # 设置垂直线条和水平线条的位置组成十字光标
                    self.vLine.setPos(mouse_point.x())
                    self.hLine.setPos(mouse_point.y())
            except Exception:
                print(traceback.print_exc())


# K线图绘制类
class CandlestickItem(pg.GraphicsObject):
    def __init__(self, data):
        pg.GraphicsObject.__init__(self)
        self.picture = QtGui.QPicture()  # 实例化一个绘图设备
        self.data = data  # data里面必须有以下字段: 时间, 开盘价, 收盘价, 最低价, 最高价
        self.generatePicture()

    def generatePicture(self):
        p = QtGui.QPainter(self.picture)  # 在picture上实例化QPainter用于绘图
        p.setPen(pg.mkPen('w'))  # 设置画笔颜色
        w = (self.data[1][0] - self.data[0][0]) / 3.
        for (t, _open, close, _min, _max) in self.data:
            # print(t, open, close, min, max)
            p.drawLine(QtCore.QPointF(t, _min), QtCore.QPointF(t, _max))  # 绘制线条
            if _open > close:  # 开盘价大于收盘价
                p.setBrush(pg.mkBrush('#45b97c'))  # 设置画刷颜色为绿
            elif _open < close:
                p.setBrush(pg.mkBrush('#ed1941'))  # 设置画刷颜色为红
            else:
                continue
            p.drawRect(QtCore.QRectF(t - w, _open, w * 2, close - _open))  # 绘制箱子
        p.end()

    def paint(self, p, *args):
        p.drawPicture(0, 0, self.picture)

    def boundingRect(self):
        return QtCore.QRectF(self.picture.boundingRect())
